n_fold <- 4
n_batch_each_epoch <- 30L
n_epoch <- 10L

library(data.table)
library(magrittr)
library(xgboost)
library(glmnet)
library(lme4)

# Read data

if (file.exists('train_aug.csv') & file.exists('test_aug.csv')) {
  
  train_dat <- fread('train_aug.csv', data.table = FALSE)
  test_dat <- fread('test_aug.csv', data.table = FALSE)
  submit_dat <- fread('sample_submission.csv', data.table = FALSE)
  
  
} else {
  
  train_dat <- fread('train.csv', data.table = FALSE)
  test_dat <- fread('test.csv', data.table = FALSE)
  submit_dat <- fread('sample_submission.csv', data.table = FALSE)
  
  # Impute
  
  interest_var <- colnames(train_dat)[-1:-3]
  
  train_dat[,'GENDER'] <- as.factor(train_dat[,'GENDER'])
  test_dat[,'GENDER'] <- as.factor(test_dat[,'GENDER'])
  
  for (i in 1:length(interest_var)) {
    if (class(train_dat[,interest_var[i]])[1] %in% c('numeric', 'integer')) {
      m_val <- mean(train_dat[,interest_var[i]], na.rm = TRUE)
    } else {
      m_val <- names(which.max(table(train_dat[,interest_var[i]])))
    }
    train_dat[train_dat[,interest_var[i]] %in% NA, interest_var[i]] <- m_val
    test_dat[test_dat[,interest_var[i]] %in% NA, interest_var[i]] <- m_val
  }
  
  # Merge more data
  
  set.seed(0)
  
  mean_y <- mean(train_dat[,'K'])
  
  pb <- txtProgressBar(max = nrow(train_dat), style = 3)
  
  for (i in 1:nrow(train_dat)) {
    
    sub_train_dat <- train_dat[train_dat[,'PID'] %in% train_dat[i,'PID'] & train_dat[,'UID'] < train_dat[i,'UID'],]
    
    if (nrow(sub_train_dat) > 0) {
      
      train_dat[i,'mean'] <- mean(sub_train_dat[,'K'])
      train_dat[i,'min'] <- min(sub_train_dat[,'K'])
      train_dat[i,'max'] <- max(sub_train_dat[,'K'])
      
      
    } else {
      
      train_dat[i,'mean'] <- mean_y
      train_dat[i,'min'] <- mean_y
      train_dat[i,'max'] <- mean_y
      
    }
    
    setTxtProgressBar(pb, i)
    
  }
  
  close(pb)
  
  pb <- txtProgressBar(max = nrow(test_dat), style = 3)
  
  for (i in 1:nrow(test_dat)) {
    
    sub_train_dat <- train_dat[train_dat[,'PID'] %in% test_dat[i,'PID'] & train_dat[,'UID'] < test_dat[i,'UID'],]
    
    if (nrow(sub_train_dat) > 0) {
      
      test_dat[i,'mean'] <- mean(sub_train_dat[,'K'])
      test_dat[i,'min'] <- min(sub_train_dat[,'K'])
      test_dat[i,'max'] <- max(sub_train_dat[,'K'])
      
    } else {
      
      test_dat[i,'mean'] <- mean_y
      test_dat[i,'min'] <- mean_y
      test_dat[i,'max'] <- mean_y
      
    }
    
    setTxtProgressBar(pb, i)
    
  }
  
  close(pb)
  
  write.csv(train_dat, 'train_aug.csv', row.names = FALSE, na = '', quote = TRUE, fileEncoding = 'CP950')
  write.csv(test_dat, 'test_aug.csv', row.names = FALSE, na = '', quote = TRUE, fileEncoding = 'CP950')
  
}

# Data aug

train_X <- model.matrix(~ ., data = train_dat[,-1:-3])
train_Y <- train_dat[,'K',drop=TRUE] %>% as.numeric()

# Get data

dtrain <- xgb.DMatrix(train_X, label = train_Y)

# Search parameter

random_param_list <- list()

random_param_list[['eta']] <- list(type = 'numeric', min = 0.05, max = 0.5, hard_min = 0.01, hard_max = 1, log = TRUE)
random_param_list[['gamma']] <- list(type = 'numeric', min = exp(-9), max = exp(2), hard_min = 1e-8, hard_max = Inf, log = TRUE)
random_param_list[['max_depth']] <- list(type = 'integer', min = 2, max = 8, hard_min = 2, hard_max = 10, log = FALSE)
random_param_list[['min_child_weight']] <- list(type = 'numeric', min = 2^0, max = 2^10, hard_min = 1, hard_max = nrow(train_X) / 4, log = TRUE)
random_param_list[['subsample']] <- list(type = 'numeric', min = 0.5, max = 1, hard_min = 0.1, hard_max = 1, log = FALSE)
random_param_list[['colsample_bytree']] <- list(type = 'numeric', min = 0.5, max = 1, hard_min = 0.1, hard_max = 1, log = FALSE)
random_param_list[['alpha']] <- list(type = 'numeric', min = 0, max = 1, hard_min = 0, hard_max = 1, log = FALSE)
random_param_list[['lambda']] <- list(type = 'numeric', min = exp(-9), max = exp(-1), hard_min = 1e-8, hard_max = 1, log = TRUE)

param_list <- list()

for (epoch in 1:n_epoch) {
  
  # Random search
  
  t1 <- Sys.time()
  
  message('Start to conduct epoch = ', epoch)
  
  for (batch in 1:n_batch_each_epoch) {
    
    ## Sampling a random params
    
    set.seed(batch + n_batch_each_epoch * epoch)
    
    current_params <- list()
    
    for (i in 1:length(random_param_list)) {
      
      min_val <- random_param_list[[i]][['min']]
      max_val <- random_param_list[[i]][['max']]
      
      if (random_param_list[[i]][['log']] %in% TRUE) {
        
        random_val <- runif(n = 1, min = log(min_val), max = log(max_val)) %>% exp()
        
      } else {
        
        random_val <- runif(n = 1, min = min_val, max = max_val)
      
      }
      
      if (random_param_list[[i]][['type']] %in% 'integer') {
        
        random_val <- random_val %>% round()
        
      }
      
      current_params[[names(random_param_list)[i]]] <- random_val
      
    }
    
    ## Training
    
    set.seed(0)
    
    cv_fit <- xgb.cv(data = dtrain, params = current_params, nfold = n_fold,
                     booster = "gbtree", metrics = "rmse", objective = "reg:linear",
                     verbose = FALSE, stratified = TRUE, maximize = FALSE,
                     nrounds = 300, early_stopping_rounds = 10, nthread = 5)
    
    evaluation_log <- cv_fit[['evaluation_log']] %>% as.data.frame()
    
    param_list[[length(param_list) + 1]] <- data.frame(epoch = epoch, 
                                                       eta = current_params[['eta']], gamma = current_params[['gamma']], 
                                                       max_depth = current_params[['max_depth']], min_child_weight = current_params[['min_child_weight']],
                                                       subsample = current_params[['subsample']], colsample_bytree = current_params[['colsample_bytree']],
                                                       alpha = current_params[['alpha']], lambda = current_params[['lambda']],
                                                       best_iter = cv_fit[['best_iteration']],
                                                       rmse = evaluation_log[cv_fit[['best_iteration']],'test_rmse_mean'])
    
  }
  
  message('time = ', formatC(as.numeric(Sys.time() - t1, unit = 'mins'), format = 'f', 1), ' mins')
  
  current_param <- do.call('rbind', param_list)
  
  # Re-set limit of random_param_list
  
  current_param <- current_param[order(current_param[,'rmse'], decreasing = FALSE),]
  current_param <- current_param[1:(n_batch_each_epoch / 2 + epoch * n_batch_each_epoch / 10),]
  current_param[,'w'] <- (current_param[,'rmse'] - min(current_param[,'rmse']) + diff(range(current_param[,'rmse'])) * 0.01)^-1
  
  for (i in 1:length(random_param_list)) {
    
    current_val <- current_param[,names(random_param_list)[i]]
    current_w <- current_param[,'w']
    
    if (random_param_list[[i]][['log']] %in% TRUE) {current_val <- log(current_val)}
    
    mean_val <- as.numeric((current_val %*% current_w) / sum(current_w))
    sd_val <- sqrt(((current_val - mean_val)^2 %*% current_w) / sum(current_w))
    best_val <- current_val[1]
    
    min_val <- mean_val - sd_val / epoch
    max_val <- mean_val + sd_val / epoch
    
    if (random_param_list[[i]][['log']] %in% TRUE) {min_val <- exp(min_val); max_val <- exp(max_val)}
    
    if (min_val < random_param_list[[i]][['hard_min']]) {min_val <- random_param_list[[i]][['hard_min']]}
    if (max_val > random_param_list[[i]][['hard_max']]) {max_val <- random_param_list[[i]][['hard_max']]}
    
    random_param_list[[i]][['min']] <- min_val
    random_param_list[[i]][['max']] <- max_val
    
  }
  
  # Print best result
  
  best_param_pos <- which.min(current_param[,'rmse'])
  print(current_param[best_param_pos,])
  
}

# Get best model

param_data <- do.call('rbind', param_list)
best_param_pos <- which.min(param_data[,'rmse'])

# tapply(param_data[,'rmse'], param_data[,'epoch'], mean) %>% plot(type = 'l')

# tapply(param_data[,'subsample'], param_data[,'epoch'], mean) %>% plot(type = 'l')
# tapply(param_data[,'subsample'], param_data[,'epoch'], sd) %>% plot(type = 'l')
# plot(param_data[,'subsample'], type = 'l')

best_model <- xgb.train(data = dtrain, nfold = n_fold,
                        booster = "gbtree", metrics = "rmse", objective = "reg:linear",
                        verbose = FALSE, stratified = TRUE, maximize = FALSE,
                        nrounds = param_data[best_param_pos,'best_iter'], nthread = 5,
                        eta = param_data[best_param_pos,'eta'], gamma = param_data[best_param_pos,'gamma'],
                        max_depth = param_data[best_param_pos,'max_depth'], min_child_weight = param_data[best_param_pos,'min_child_weight'],
                        subsample = param_data[best_param_pos,'subsample'], colsample_bytree = param_data[best_param_pos,'colsample_bytree'],
                        alpha = param_data[best_param_pos,'alpha'], lambda = param_data[best_param_pos,'lambda'])
  
# Get prediction

train_dat[,'xgb_pred'] <- predict(best_model, train_X)

test_X <- model.matrix(~ ., data = test_dat[,-1:-2])
test_dat[,'xgb_pred'] <- predict(best_model, test_X)

# Revise data  

n_PID <- table(train_dat[,'PID'])

sub_train_dat <- train_dat[train_dat[,'PID'] %in% names(n_PID)[n_PID > 1],]

lmm_fit <- lmer(K ~ xgb_pred + (1 + xgb_pred | PID), data = sub_train_dat)

G_mat <- matrix(as.numeric(VarCorr(lmm_fit)[[1]]), nrow = 2, ncol = 2)

pb <- txtProgressBar(max = nrow(test_dat), style = 3)

for (i in 1:nrow(test_dat)) {
  
  Z_mat <- cbind(1, train_dat[train_dat[,'PID'] %in% test_dat[i,'PID'], 'xgb_pred'])
  R_mat <- matrix(0, nrow = nrow(Z_mat), ncol = nrow(Z_mat))
  diag(R_mat) <- sigma(lmm_fit)^2
  Sigma_mat <- Z_mat %*% G_mat %*% t(Z_mat) + R_mat
  Res_val <- train_dat[train_dat[,'PID'] %in% test_dat[i,'PID'], 'K'] - Z_mat %*% fixef(lmm_fit)
  BLUP <- G_mat %*% t(Z_mat) %*% solve(Sigma_mat) %*% Res_val
  person_coef <- fixef(lmm_fit) + BLUP
  test_dat[i,'b0_person'] <- person_coef[1]
  test_dat[i,'b1_person'] <- person_coef[2]
  
  setTxtProgressBar(pb, i)
  
}

close(pb)

# Write out

submit_dat[,'K'] <- test_dat[,'b0_person'] + test_dat[,'b1_person'] * test_dat[,'xgb_pred']
fwrite(submit_dat, file = 'my_submission.csv', na = '', row.names = FALSE, quote = FALSE)